import json
import os
from zhipuai import ZhipuAI
from dotenv import load_dotenv, find_dotenv

_ = load_dotenv(find_dotenv())
api_key = os.environ.get('ZHIPUAI_API_KEY')
if api_key is None:
    raise ValueError("API Key is not set in the .env file")
client = ZhipuAI(api_key=api_key)

import numpy as np  
from scipy.integrate import quad  

# 定义被积函数  f(x) = sin(x)
def f(x):  
    return np.sin(x) 

# 计算定积分
def di(x, y):
    return quad(f, x, y) [0]

# 创建一个用于Function Calling的函数定义
def create_function_calling_request(messages):
    return client.chat.completions.create(
        model="glm-4-plus",
        messages=messages,
        tools=[
            {
                "type": "function",
                "function": {
                    "name": "di",
                    "description": "计算函数f(x) = sin(x)的在[x,y]的定积分。",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "x": {
                                "type": "number", 
                                "description": """定积分下限，
                                必须从用户输入中识别出定积分下限，
                                否则提示用户重新输入。"""
                            }, 
                            "y": {
                                "type": "number", 
                                "description": """定积分上限，
                                必须从用户输入中识别出定积分下限，
                                否则提示用户重新输入。"""
                            }
                        },
                        "required": ["x", "y"]
                    }
                }
            }
        ],
        tool_choice="auto"
    )

# 创建一个包含函数调用的prompt
prompt = "计算函数f(x) = sin(x)的在[0,π]的定积分。"
messages=[{"role": "user", "content": prompt}]

# 发送请求并获取响应
response = create_function_calling_request(messages)
print(f"===Debug===\n{response}\n")

# 解析响应并调用相应的函数
function_call = response.choices[0].message.tool_calls[0]
if function_call is  None:
    result = response.choices[0].message.content.strip()
else:
    messages.append(response.choices[0].message.model_dump())
    function_name = function_call.function.name
    function_args = json.loads(function_call.function.arguments)

    if function_name == "di":
        a = function_args.get("x")
        b = function_args.get("y")
        result = di(a, b)
        
        # 将本地计算结果返回大模型
        follow_up_prompt = {
            "role": "tool",
            "tool_call_id": function_call.id,
            "content": json.dumps({"result": result})
        }
        messages.append(follow_up_prompt)
        
        print(f"===Debug===\nmessages=\n{messages}\n")
        # 使用tool_call.id向大模型返回结果
        final_response = create_function_calling_request(messages)
        print(f"===Debug===\n{final_response}\n")

        # 最终结果
        result = final_response.choices[0].message.content.strip()
    else:
        print("Function not recognized.")

# 输出最终结果
print(result)
  